

#include <string.h>
#include "XMen_sm3.h"
#include "XMen_asn1.h"


/**
 * @brief Rotate left operation
 * @param x Value to rotate
 * @param n Number of bits to rotate by
 * @return Rotated value
 */
#define ROTL(x, n) (((x) << (n)) | ((x) >> (32 - (n))))


/**
 * @brief Permutation function P0 used in message expansion
 * @param x Input value
 * @return P0(x) = x XOR (x rotated left 9) XOR (x rotated left 17)
 */
#define P0(x) ((x) ^ ROL32((x), 9) ^ ROL32((x), 17))


/**
 * @brief Permutation function P1 used in message expansion
 * @param x Input value
 * @return P1(x) = x XOR (x rotated left 15) XOR (x rotated left 23)
 */
#define P1(x) ((x) ^ ROL32((x), 15) ^ ROL32((x), 23))


/**
 * @brief Boolean function FF00 used in rounds 0-15
 * @param x First input
 * @param y Second input
 * @param z Third input
 * @return FF00(x,y,z) = x XOR y XOR z
 */
#define FF00(x, y, z) ((x) ^ (y) ^ (z))


/**
 * @brief Boolean function FF16 used in rounds 16-63
 * @param x First input
 * @param y Second input
 * @param z Third input
 * @return FF16(x,y,z) = (x AND y) OR (x AND z) OR (y AND z)
 */
#define FF16(x, y, z) (((x) & (y)) | ((x) & (z)) | ((y) & (z)))


/**
 * @brief Boolean function GG00 used in rounds 0-15
 * @param x First input
 * @param y Second input
 * @param z Third input
 * @return GG00(x,y,z) = x XOR y XOR z
 */
#define GG00(x, y, z) ((x) ^ (y) ^ (z))


/**
 * @brief Boolean function GG16 used in rounds 16-63
 * @param x First input
 * @param y Second input
 * @param z Third input
 * @return GG16(x,y,z) = (x AND (y XOR z)) XOR z
 */
#define GG16(x, y, z) ((((y) ^ (z)) & (x)) ^ (z))


/**
 * @brief Single round of SM3 compression function
 * Performs one round of the compression function updating working variables
 */
#define R(A, B, C, D, E, F, G, H, xx)                    \
	SM3_SS1 = ROL32((ROL32(A, 12) + E + ARRAY_K[j]), 7);           \
	SM3_SS2 = SM3_SS1 ^ ROL32(A, 12);                            \
	SM3_TT1 = FF##xx(A, B, C) + D + SM3_SS2 + (W[j] ^ W[j + 4]); \
	SM3_TT2 = GG##xx(E, F, G) + H + SM3_SS1 + W[j];              \
	B = ROL32(B, 9);                                     \
	H = SM3_TT1;                                             \
	F = ROL32(F, 19);                                    \
	D = P0(SM3_TT2);                                         \
	j++


/**
 * @brief Eight rounds of SM3 compression function
 * Performs 8 rounds of compression updating all working variables
 */
#define Round8(A, B, C, D, E, F, G, H, xx) \
	R(A, B, C, D, E, F, G, H, xx);     \
	R(H, A, B, C, D, E, F, G, xx);     \
	R(G, H, A, B, C, D, E, F, xx);     \
	R(F, G, H, A, B, C, D, E, xx);     \
	R(E, F, G, H, A, B, C, D, xx);     \
	R(D, E, F, G, H, A, B, C, xx);     \
	R(C, D, E, F, G, H, A, B, xx);     \
	R(B, C, D, E, F, G, H, A, xx)


/**
 * @brief Constants used in SM3 compression function
 * These are the round constants Tj defined in the SM3 specification
 * For j=0 to 15: Tj = 0x79cc4519
 * For j=16 to 63: Tj = 0x7a879d8a
 */
#define value_K0 0x79cc4519U
#define value_K1 0xf3988a32U
#define value_K2 0xe7311465U
#define value_K3 0xce6228cbU
#define value_K4 0x9cc45197U
#define value_K5 0x3988a32fU
#define value_K6 0x7311465eU
#define value_K7 0xe6228cbcU
#define value_K8 0xcc451979U
#define value_K9 0x988a32f3U
#define value_K10 0x311465e7U
#define value_K11 0x6228cbceU
#define value_K12 0xc451979cU
#define value_K13 0x88a32f39U
#define value_K14 0x11465e73U
#define value_K15 0x228cbce6U
#define value_K16 0x9d8a7a87U
#define value_K17 0x3b14f50fU
#define value_K18 0x7629ea1eU
#define value_K19 0xec53d43cU
#define value_K20 0xd8a7a879U
#define value_K21 0xb14f50f3U
#define value_K22 0x629ea1e7U
#define value_K23 0xc53d43ceU
#define value_K24 0x8a7a879dU
#define value_K25 0x14f50f3bU
#define value_K26 0x29ea1e76U
#define value_K27 0x53d43cecU
#define value_K28 0xa7a879d8U
#define value_K29 0x4f50f3b1U
#define value_K30 0x9ea1e762U
#define value_K31 0x3d43cec5U
#define value_K32 0x7a879d8aU
#define value_K33 0xf50f3b14U
#define value_K34 0xea1e7629U
#define value_K35 0xd43cec53U
#define value_K36 0xa879d8a7U
#define value_K37 0x50f3b14fU
#define value_K38 0xa1e7629eU
#define value_K39 0x43cec53dU
#define value_K40 0x879d8a7aU
#define value_K41 0x0f3b14f5U
#define value_K42 0x1e7629eaU
#define value_K43 0x3cec53d4U
#define value_K44 0x79d8a7a8U
#define value_K45 0xf3b14f50U
#define value_K46 0xe7629ea1U
#define value_K47 0xcec53d43U
#define value_K48 0x9d8a7a87U
#define value_K49 0x3b14f50fU
#define value_K50 0x7629ea1eU
#define value_K51 0xec53d43cU
#define value_K52 0xd8a7a879U
#define value_K53 0xb14f50f3U
#define value_K54 0x629ea1e7U
#define value_K55 0xc53d43ceU
#define value_K56 0x8a7a879dU
#define value_K57 0x14f50f3bU
#define value_K58 0x29ea1e76U
#define value_K59 0x53d43cecU
#define value_K60 0xa7a879d8U
#define value_K61 0x4f50f3b1U
#define value_K62 0x9ea1e762U
#define value_K63 0x3d43cec5U


/**
 * @brief Array of round constants used in compression function
 * Contains the 64 constants Tj used in each round
 */
static uint32_t ARRAY_K[64] = {
    value_K0, value_K1, value_K2, value_K3, value_K4, value_K5, value_K6, value_K7,
    value_K8, value_K9, value_K10, value_K11, value_K12, value_K13, value_K14, value_K15,
    value_K16, value_K17, value_K18, value_K19, value_K20, value_K21, value_K22, value_K23,
    value_K24, value_K25, value_K26, value_K27, value_K28, value_K29, value_K30, value_K31,
    value_K32, value_K33, value_K34, value_K35, value_K36, value_K37, value_K38, value_K39,
    value_K40, value_K41, value_K42, value_K43, value_K44, value_K45, value_K46, value_K47,
    value_K48, value_K49, value_K50, value_K51, value_K52, value_K53, value_K54, value_K55,
    value_K56, value_K57, value_K58, value_K59, value_K60, value_K61, value_K62, value_K63,
};


/**
 * @brief Compress message blocks in SM3 hash function
 * @param digest Current hash value (8 words)
 * @param data Input message blocks
 * @param blocks Number of 512-bit blocks to process
 *
 * This function implements the compression function of SM3.
 * It processes message blocks and updates the intermediate hash value.
 */
void XMen_lib_sm3_compress_blocks(uint32_t digest[8], const uint8_t* data, uint32_t blocks)
{
    uint32_t temp_A;
    uint32_t temp_B;
    uint32_t temp_C;
    uint32_t temp_D;
    uint32_t temp_E;
    uint32_t temp_F;
    uint32_t temp_G;
    uint32_t temp_H;
    uint32_t W[68];
    uint32_t SM3_SS1, SM3_SS2, SM3_TT1, SM3_TT2;
    int j;
    while (blocks--)
    {
        // Load current hash value into working variables
        temp_A = digest[0];
        temp_B = digest[1];
        temp_C = digest[2];
        temp_D = digest[3];
        temp_E = digest[4];
        temp_F = digest[5];
        temp_G = digest[6];
        temp_H = digest[7];

        // Fill first 16 words of message schedule from input block
        for (j = 0; j < 16; j++)
            W[j] = GETU32(data + j * 4);

        // Expand message schedule for remaining rounds
        for (; j < 68; j++)
            W[j] = P1(W[j - 16] ^ W[j - 9] ^ ROL32(W[j - 3], 15)) ^ ROL32(W[j - 13], 7) ^ W[j - 6];

        j = 0;

        // 64 rounds of compression
        Round8(temp_A, temp_B, temp_C, temp_D, temp_E, temp_F, temp_G, temp_H, 00);
        Round8(temp_A, temp_B, temp_C, temp_D, temp_E, temp_F, temp_G, temp_H, 00);
        Round8(temp_A, temp_B, temp_C, temp_D, temp_E, temp_F, temp_G, temp_H, 16);
        Round8(temp_A, temp_B, temp_C, temp_D, temp_E, temp_F, temp_G, temp_H, 16);
        Round8(temp_A, temp_B, temp_C, temp_D, temp_E, temp_F, temp_G, temp_H, 16);
        Round8(temp_A, temp_B, temp_C, temp_D, temp_E, temp_F, temp_G, temp_H, 16);
        Round8(temp_A, temp_B, temp_C, temp_D, temp_E, temp_F, temp_G, temp_H, 16);
        Round8(temp_A, temp_B, temp_C, temp_D, temp_E, temp_F, temp_G, temp_H, 16);

        // Update hash value
        digest[0] ^= temp_A;
        digest[1] ^= temp_B;
        digest[2] ^= temp_C;
        digest[3] ^= temp_D;
        digest[4] ^= temp_E;
        digest[5] ^= temp_F;
        digest[6] ^= temp_G;
        digest[7] ^= temp_H;

        // Move to next block
        data += 64;
    }
}


/**
 * @brief Initialize SM3 hash context
 * @param ctx Context to initialize
 *
 * Sets up initial hash value and zeroes other fields
 */
void XMen_lib_sm3_init(LIB_SM3_CTX* ctx)
{
    memset(ctx, 0, sizeof(*ctx));
    ctx->digest[0] = 0x7380166Fu;
    ctx->digest[1] = 0x4914B2B9u;
    ctx->digest[2] = 0x172442D7u;
    ctx->digest[3] = 0xDA8A0600u;
    ctx->digest[4] = 0xA96F30BCu;
    ctx->digest[5] = 0x163138AAu;
    ctx->digest[6] = 0xE38DEE4Du;
    ctx->digest[7] = 0xB0FB0E4Eu;
}


/**
 * @brief Update SM3 hash with input data
 * @param ctx Hash context
 * @param data Input data buffer
 * @param data_len Length of input in bytes
 *
 * Processes input data and updates hash state.
 * Handles partial blocks and buffering as needed.
 */
void XMen_lib_sm3_update(LIB_SM3_CTX* ctx, const uint8_t* data, uint32_t data_len)
{
    uint32_t blocks;
    
    ctx->num &= 0x3f;

    // Handle data remaining from previous update
    if (ctx->num && data_len > 0)
    {
        uint32_t left = LIB_SM3_BLOCK_SIZE - ctx->num;
        uint32_t copy_len = (data_len < left) ? data_len : left;
        memcpy(ctx->block + ctx->num, data, copy_len);
        ctx->num += copy_len;
        data += copy_len;
        data_len -= copy_len;

        // Process block if it's full
        if (ctx->num == LIB_SM3_BLOCK_SIZE)
        {
            XMen_lib_sm3_compress_blocks(ctx->digest, ctx->block, 1);
            ctx->nblocks++;
            ctx->num = 0;
        }
    }

    // Process full blocks directly from input
     blocks = data_len / LIB_SM3_BLOCK_SIZE;
    if (blocks > 0)
    {
        XMen_lib_sm3_compress_blocks(ctx->digest, data, blocks);
        ctx->nblocks += blocks;
        data += LIB_SM3_BLOCK_SIZE * blocks;
        data_len -= LIB_SM3_BLOCK_SIZE * blocks;
    }

    // Store remaining data for next update
    if (data_len > 0)
    {
        memcpy(ctx->block, data, data_len);
        ctx->num = data_len;
    }
}


/**
 * @brief Finalize SM3 hash computation
 * @param ctx Hash context
 * @param digest Output buffer for hash value (32 bytes)
 *
 * Performs padding, processes final block(s) and outputs hash value.
 * Follows SM3 padding rules: append 1 bit, then zeros, then 64-bit length.
 */
void XMen_lib_sm3_finish(LIB_SM3_CTX* ctx, uint8_t* digest)
{
    int i;

    // Add padding bit
    ctx->num &= 0x3f;
    ctx->block[ctx->num] = 0x80;

    // Calculate padding length
    size_t padding_length = (ctx->num <= LIB_SM3_BLOCK_SIZE - 9) ? LIB_SM3_BLOCK_SIZE - ctx->num - 9 : LIB_SM3_BLOCK_SIZE - ctx->num - 1;

    // Add padding zeros
    memset(ctx->block + ctx->num + 1, 0, padding_length);

    // Handle case where length doesn't fit in current block
    if (ctx->num > LIB_SM3_BLOCK_SIZE - 9)
    {
        XMen_lib_sm3_compress_blocks(ctx->digest, ctx->block, 1);
        memset(ctx->block, 0, LIB_SM3_BLOCK_SIZE - 8);
    }

    // Add 64-bit message length
    PUTU32(ctx->block + 56, ctx->nblocks >> 23);
    PUTU32(ctx->block + 60, (ctx->nblocks << 9) + (ctx->num << 3));

    // Process final block and output hash value
    XMen_lib_sm3_compress_blocks(ctx->digest, ctx->block, 1);
    for (i = 0; i < 8; i++)
    {
        PUTU32(digest + i * 4, ctx->digest[i]);
    }
}


/**
 * @brief Compute SM3 hash of message in one call
 * @param msg Input message buffer
 * @param msglen Length of message in bytes
 * @param dgst Output buffer for hash value (32 bytes)
 *
 * Convenience function that handles initialization,
 * update and finalization in one call.
 */
void XMen_lib_sm3_digest(const uint8_t* msg, uint32_t msglen, uint8_t dgst[LIB_SM3_DIGEST_SIZE])
{
    LIB_SM3_CTX ctx;
    XMen_lib_sm3_init(&ctx);
    XMen_lib_sm3_update(&ctx, msg, msglen);
    XMen_lib_sm3_finish(&ctx, dgst);
    memset(&ctx, 0, sizeof(ctx));
}